# careful: this overwrites the built-in abs function
from numpy import int16,abs,where,array,concatenate,delete

Ckms = 299792.458         # speed of light km/s, exact

#@profile
def find_groups(zgal,za,zb,zc,vdiff=None):#,multgal=False):
    """ Given a list of galaxy redshifts and absorber redshifts in
    three nearby sight-lines, for each galaxy finds any absorbers
    within some given velocity range around that galaxy.

    Parameters
    ----------
    Galaxy redshift array, and absorber redshift arrays for each LOS.
    Also an optional keyword parameter giving the maximum velocity
    difference between the galaxy and absorbers for the absorbers to
    be considered 'associated' with the galaxy.

    new option: if multgal=True, then as well as searching for
    matching absorbers within the velocity range, we will search for
    other matching galaxies in this velocity range. Thus there can be
    multiple galaxies and absorbers per group. After all groups have
    been measured, pass through the list and remove any groups whose
    components (each galaxy and absorber that makes up that group) are
    all contained within another single group.

    Returns
    -------
    A list with one entry for each galaxy. Each entry is a dictionary
    (one key per LOS) giving the indices of any absorbers that have <=
    (velocity difference) from the galaxy.
    """

    # seems to take up too much memory recording i and dv.  just
    # record i for now.
    #dtype = [('i','i2'),('dv','f4')]
    groups = []
    for z in zgal:
        # calculate v difference with every absorber for each
        # sight-line.  Use approximate expression for dv, since we're
        # dealing with velocities < 0.01c
        indices = dict()
        for z1,key in zip([za,zb,zc],'abc'):
            dv = (z - z1) / (1 + (z + z1)/2.) * Ckms
            # select those absorbers with absolute vel difference < vdiff
            ind, = where(abs(dv) <= vdiff)
            indices[key] = zip(ind.astype(int16),dv[ind])
            #indices[key] = ind.astype(int16)
            
#         if multgal:
#             # Galaxies
#             dv = (z - zgal) / (1 + (z + zgal)/2.) * Ckms
#             ind, = where(abs(dv) <= vdiff)
#             #dv = dv[ind]
#             indices['g'] = ind.astype(int16)

        groups.append(indices)

#     if multgal:
#         # select only 'unique' groups (as defined in function docstring)
#         ugroups = []
#         for g in groups:
#             for u in ugroups:
#                 cond = (set(g['a']) - set(u['a']),
#                         set(g['b']) - set(u['b']),
#                         set(g['c']) - set(u['c']),
#                         set(g['g']) - set(u['g']))
#                 if not any(cond):
#                     break
#             else:
#                 ugroups.append(g)
#         groups = ugroups
    
    return tuple(groups)

#@profile
def find_abs_groups(za,zb,zc,vdiff=None):
    """ For each absorber redshift, finds unique groups of
    'associated' absorbers across the three sight-lines within some
    velocity range around that redshift.

    Here 'unique' means that for all the absorbers in any given group,
    there is no other group that also contains all of those absorbers.

    Parameters
    ----------
    za,zb,zc:
        Absorber redshift arrays for each LOS.
        
    vdiff=None:
        Search this far in velocity space to look for an associated
        absorber.

    Returns
    -------
    list of groups:
        The indices of each line in the group and the maximum v
        difference between lines in that group.
    """

    allz = concatenate([za,zb,zc])
    groups = []
    for z0 in sorted(allz):
        indices = dict()
        zall = []
        for key,z in zip('abc',[za,zb,zc]):
            # select absorbers which are within vdiff of the reference
            # absorber redshift
            dv = Ckms * (z0-z) / (1 + (z0+z)/2.)
            abs_dv = abs(dv)
            ind, = (abs_dv <= vdiff).nonzero()
            # useful to store as a set for remove_ununique
            indices[key] = set(ind)
            zall.extend(list(z[ind]))

        # record the largest velocity difference between absorbers in
        # this group (will be 0 when only one absorber present)
        zall = array(zall)
        dv = (zall.max() - zall)/(1 + zall.mean())*Ckms
        indices['dvmax'] = dv.max()
        groups.append(indices)

    # now only keep 'unique' groups.
#     ugroups = []
#     for g in groups:
#         for u in ugroups:
#             cond = (set(g['a']) - set(u['a']),
#                     set(g['b']) - set(u['b']),
#                     set(g['c']) - set(u['c']))
#             if not any(cond):
#                 break
#         else:
#             ugroups.append(g)

    return remove_ununique(groups)

def remove_ununique(groups):
    """ Removes groups that are not 'unique'."""

    ugroups0 = []
    # this gets rid of most of the duplicate groups
    for g in groups:
        for u in ugroups0:
            # set differences
            cond = (g['a'] - u['a'], g['b'] - u['b'], g['c'] - u['c'])
            if not any(cond):
#                 print 'g (not adding)',g
#                 print 'u (comparison)',u
#                 raw_input('enter')
                break
        else:
            ugroups0.append(g)
    #print len(groups),len(ugroups0)
    # need to run this again to remove any groups that are not
    # identical, but one still could be wholly contained within
    # another.
    ugroups1 = []
    for i0,g0 in enumerate(ugroups0):
        for i1,g1 in enumerate(ugroups0):
            if i1 == i0:  continue
            # set differences
            cond = (g0['a'] - g1['a'], g0['b'] - g1['b'], g0['c'] - g1['c'])
            if not any(cond):
#                 print 'g0 (removing)',g0
#                 print 'g1 (comparison)',g1
#                 raw_input('enter')
                break
        else:
            ugroups1.append(g0)
    #print len(ugroups0),len(ugroups1)

    return ugroups1

def find_unique_absgroups(groupabs):
    """ Given the output of the above routine, selects groups with
    absorbers across two or three sight-lines.

    Parameters
    ----------
    groupabs:
        group of input absorbers (output of find_abs_groups())
    
    Returns
    -------
    triples,doubles:
        Two lists of indices. One list entry for each double/triple
        group. A triple is a group where at least one line appears in
        all three sight-lines. A double is group where at least one
        line appears in at least two sight-lines.  Note triples are
        also counted as doubles.
    """

    doubles = []
    triples = []

    for g in groupabs:
        lena,lenb,lenc = [len(g[key]) for key in 'abc']
        if all([lena,lenb,lenc]):
            triples.append(g)
            doubles.append(g)
        elif (lena and lenb) or (lenb and lenc) or (lena and lenc):
            doubles.append(g)
            
    return triples,doubles


def nabs_group(group):
    """ Find number of galaxies which have (00) No associated
    absorbers, (0) associated absorbers in only one line of sight, (1)
    one or more associated absorber (in one, two or three sightlines),
    (2) two or three sight lines with at least 1 associated absorber,
    (3) three sight lines with at least 1 associated absorber."""
    
    n00 = 0;  n0 = 0;  n1 = 0;  n2 = 0;  n3 = 0
    for g in group:
        lena,lenb,lenc = [len(g[key]) for key in 'abc']
        if all([lena,lenb,lenc]):
            n3 += 1
            n2 += 1
            n1 += 1
        elif (lena and lenb) or (lena and lenc) or (lenb and lenc):
            n2 += 1
            n1 += 1
        elif any([lena,lenb,lenc]):
            n1 += 1
            n0 += 1
        else:
            n00 += 1

    return  n0, n1, n2, n3
